package com.md.infrastructure.configuration;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.md.infrastructure.config.AzureSamlConfig;
import com.md.util.ClassUtils;
import com.md.util.CollectionUtil;
import com.md.util.ObjectUtil;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import net.shibboleth.utilities.java.support.xml.ParserPool;
import org.opensaml.core.config.ConfigurationService;
import org.opensaml.core.xml.XMLObject;
import org.opensaml.core.xml.config.XMLObjectProviderRegistry;
import org.opensaml.core.xml.io.Unmarshaller;
import org.opensaml.saml.common.xml.SAMLConstants;
import org.opensaml.saml.ext.saml2alg.SigningMethod;
import org.opensaml.saml.saml2.core.AuthnRequest;
import org.opensaml.saml.saml2.core.Issuer;
import org.opensaml.saml.saml2.core.NameID;
import org.opensaml.saml.saml2.core.NameIDPolicy;
import org.opensaml.saml.saml2.core.impl.AuthnRequestBuilder;
import org.opensaml.saml.saml2.core.impl.AuthnRequestMarshaller;
import org.opensaml.saml.saml2.core.impl.IssuerBuilder;
import org.opensaml.saml.saml2.core.impl.NameIDBuilder;
import org.opensaml.saml.saml2.core.impl.NameIDPolicyBuilder;
import org.opensaml.saml.saml2.metadata.EntitiesDescriptor;
import org.opensaml.saml.saml2.metadata.EntityDescriptor;
import org.opensaml.saml.saml2.metadata.Extensions;
import org.opensaml.saml.saml2.metadata.IDPSSODescriptor;
import org.opensaml.security.credential.UsageType;
import org.opensaml.xmlsec.keyinfo.KeyInfoSupport;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.ResourceLoader;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import javax.xml.namespace.QName;
import java.io.IOException;
import java.io.InputStream;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.List;
import java.util.Map;

/**
 * Azure ED 授权登录 核心配置
 * 使用前必须初始化链接
 *
 * @author 刘骄阳
 * 2024-03-24 3:04
 * @see AfterRunner
 */
@UtilityClass
@Slf4j
public class AzureSamlCoreConfiguration {
  public static final XMLObjectProviderRegistry REGISTRY = ConfigurationService.get(XMLObjectProviderRegistry.class);
  public static final ParserPool PARSER_POOL = REGISTRY.getParserPool();
  public static final ResourceLoader RESOURCE_LOADER = new DefaultResourceLoader();
  public static final AuthnRequestBuilder AUTHN_REQUEST_BUILDER = getRegistryClass(AuthnRequestBuilder.class,
    AuthnRequest.DEFAULT_ELEMENT_NAME);
  public static final AuthnRequestMarshaller MARSHALLER = (AuthnRequestMarshaller) REGISTRY.getMarshallerFactory().
    getMarshaller(AuthnRequest.DEFAULT_ELEMENT_NAME);

  public static final IssuerBuilder ISSUER_BUILDER = getRegistryClass(IssuerBuilder.class,
    Issuer.DEFAULT_ELEMENT_NAME);
  public static final NameIDBuilder NAME_ID_BUILDER = getRegistryClass(NameIDBuilder.class,
    NameID.DEFAULT_ELEMENT_NAME);
  public static final NameIDPolicyBuilder NAME_ID_POLICY_BUILDER = getRegistryClass(NameIDPolicyBuilder.class,
    NameIDPolicy.DEFAULT_ELEMENT_NAME);

  /**
   * key: 应用id
   * value: 应用对应的登录配置
   */
  public static final Map<String, AzureSamlConfig> AZURE_SAML_CONFIG_MAP = Maps.newHashMap();


  @SuppressWarnings("unchecked")
  private static <T> T getRegistryClass(final Class<T> clazz, final QName key) {
    if (ObjectUtil.isEmpty(REGISTRY)) {
      throw ServerException.build(ErrorEnum.GET_XML_REGISTRY_ERROR, "解析器为空");
    }
    T t = (T) REGISTRY.getBuilderFactory().getBuilder(key);
    if (ObjectUtil.isEmpty(t)) {
      throw ServerException.build(ErrorEnum.GET_XML_REGISTRY_ERROR, clazz.getName() + "为空");
    }
    return t;
  }


  public List<AzureSamlConfig> analysis(final String url) {
    try (InputStream source = AzureSamlCoreConfiguration.RESOURCE_LOADER.getResource(url).getInputStream()) {
      return analysis(source);
    } catch (IOException e) {
      throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR, e);
    }
  }

  public List<AzureSamlConfig> analysis(final InputStream inputStream) {
    List<AzureSamlConfig> builders = Lists.newArrayList();
    XMLObject xmlObject = xmlObject(inputStream);
    if (xmlObject instanceof EntitiesDescriptor) {
      EntitiesDescriptor descriptors = (EntitiesDescriptor) xmlObject;
      for (EntityDescriptor descriptor : descriptors.getEntityDescriptors()) {
        if (descriptor.getIDPSSODescriptor(SAMLConstants.SAML20P_NS) != null) {
          builders.add(getAzureSamlConfig(descriptor));
        }
      }
      if (builders.isEmpty()) {
        // 未读取到 IDPSSODescriptor 配置,IDPSSODescriptor元素包含了身份提供者用来接受和处理身份验证请求（AuthnRequest）的服务信息
        throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR);
      }
      return builders;
    }
    if (xmlObject instanceof EntityDescriptor) {
      EntityDescriptor descriptor = (EntityDescriptor) xmlObject;
      return Lists.newArrayList(getAzureSamlConfig(descriptor));
    }
    // 未知的类型, xmlObject.getClass()
    throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR);
  }

  /**
   * 读取流,读取成XMLObject
   *
   * @param inputStream 流
   * @return xmlObject
   */
  private XMLObject xmlObject(final InputStream inputStream) {
    Document document = document(inputStream);
    Element element = document.getDocumentElement();
    Unmarshaller unmarshaller = AzureSamlCoreConfiguration.REGISTRY.getUnmarshallerFactory().getUnmarshaller(element);
    if (unmarshaller == null) {
      log.error("Unsupported element of type  {}", element.getTagName());
      throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR);
    }
    try {
      return unmarshaller.unmarshall(element);
    } catch (Exception ex) {
      throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR, ex);
    }
  }


  private Document document(final InputStream inputStream) {
    try {
      return PARSER_POOL.parse(inputStream);
    } catch (Exception ex) {
      throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR, ex);
    }
  }

  /**
   * 获取saml登录配置
   *
   * @param descriptor
   * @return
   */
  private AzureSamlConfig getAzureSamlConfig(final EntityDescriptor descriptor) {
    IDPSSODescriptor idpssoDescriptor = descriptor.getIDPSSODescriptor(SAMLConstants.SAML20P_NS);
    if (idpssoDescriptor == null) {
      // 元数据不完整,未读取到所需数据
      throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR);
    }
    List<X509Certificate> verification = Lists.newArrayList();
    List<X509Certificate> encryption = Lists.newArrayList();

    idpssoDescriptor.getKeyDescriptors().forEach(keyDescriptor -> {
      try {
        List<X509Certificate> certificates = KeyInfoSupport.getCertificates(keyDescriptor.getKeyInfo());
        if (ObjectUtil.hasEquals(keyDescriptor.getUse(), UsageType.SIGNING, UsageType.UNSPECIFIED)) {
          verification.addAll(certificates);
        }
        if (ObjectUtil.hasEquals(keyDescriptor.getUse(), UsageType.ENCRYPTION, UsageType.UNSPECIFIED)) {
          encryption.addAll(certificates);
        }
      } catch (CertificateException ex) {
        throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR, ex);
      }
    });
    if (verification.isEmpty()) {
      // 未找到公钥证书
      throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR);
    }

    if (idpssoDescriptor.getSingleSignOnServices().isEmpty()) {
      // 未包含IdP接收和处理AuthnRequest的URL地址。
      throw ServerException.build(ErrorEnum.ANALYSIS_METADATA_ERROR);
    }
    AzureSamlConfig azureSamlConfig = AzureSamlConfig.builder().
      entityId(descriptor.getEntityID())
      .needSigned(Boolean.TRUE.equals(idpssoDescriptor.getWantAuthnRequestsSigned()))
      .encryptionCertificates(encryption)
      .verificationCertificates(verification)
      .descriptor(descriptor)
      .signingAlgorithms(CollectionUtil.map(signingMethods(idpssoDescriptor), SigningMethod::getAlgorithm))
      .build();
    azureSamlConfig.setSignOnServiceBinding(idpssoDescriptor.getSingleSignOnServices());
    azureSamlConfig.setSignOutServiceBinding(idpssoDescriptor.getSingleLogoutServices());
    return azureSamlConfig;
  }

  /**
   * 获取签名算法
   *
   * @param idpSsoDescriptor idp信息
   * @return 签名算法
   */
  private List<SigningMethod> signingMethods(final IDPSSODescriptor idpSsoDescriptor) {
    Extensions extensions = idpSsoDescriptor.getExtensions();
    List<SigningMethod> result = signingMethods(extensions);
    if (!result.isEmpty()) {
      return result;
    }
    EntityDescriptor descriptor = (EntityDescriptor) idpSsoDescriptor.getParent();
    if (ObjectUtil.isEmpty(descriptor)) {
      return Lists.newArrayList();
    }
    extensions = descriptor.getExtensions();
    return signingMethods(extensions);
  }

  private List<SigningMethod> signingMethods(final Extensions extensions) {
    if (extensions != null) {
      return ClassUtils.source2Target(extensions.getUnknownXMLObjects(SigningMethod.DEFAULT_ELEMENT_NAME),
        SigningMethod.class);
    }
    return Lists.newArrayList();
  }


}
